import sys
import os
import cv2
import torch
import numpy as np
from PyQt5.QtWidgets import (QApplication, QMainWindow, QVBoxLayout, QHBoxLayout, 
                             QPushButton, QLabel, QComboBox, QTextEdit, QGroupBox,
                             QWidget, QFileDialog, QMessageBox, QProgressBar,
                             QTabWidget, QCheckBox, QSpinBox, QDoubleSpinBox)
from PyQt5.QtCore import QTimer, Qt, pyqtSignal, QThread
from PyQt5.QtGui import QImage, QPixmap, QFont
from ultralytics import YOLO

class DetectionThread(QThread):
    """检测线程"""
    finished = pyqtSignal()
    result_ready = pyqtSignal(np.ndarray, list, str)
    error_occurred = pyqtSignal(str)
    
    def __init__(self, model_path, source, conf_threshold=0.5):
        super().__init__()
        self.model_path = model_path
        self.source = source
        self.conf_threshold = conf_threshold
        self.running = False
        
    def run(self):
        try:
            # 加载模型
            self.model = YOLO(self.model_path)
            self.running = True
            
            # 处理不同的输入源
            if isinstance(self.source, str) and os.path.isfile(self.source):
                # 视频文件
                cap = cv2.VideoCapture(self.source)
                while self.running and cap.isOpened():
                    ret, frame = cap.read()
                    if not ret:
                        break
                    
                    # 执行检测
                    results = self.model(frame, conf=self.conf_threshold, verbose=False)
                    
                    # 处理结果
                    detected_objects = []
                    for result in results:
                        boxes = result.boxes
                        if boxes is not None:
                            for box in boxes:
                                cls = int(box.cls.item())
                                conf = box.conf.item()
                                detected_objects.append((cls, conf))
                    
                    # 分类驾驶行为 - 支持多重行为检测
                    behavior_classes = self.classify_behavior(detected_objects)
                    
                    # 绘制检测结果
                    annotated_frame = results[0].plot()
                    
                    # 添加行为分类文本
                    try:
                        from PIL import Image, ImageDraw, ImageFont
                        import numpy as np
                        
                        # 转换OpenCV图像到PIL图像
                        pil_image = Image.fromarray(cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB))
                        draw = ImageDraw.Draw(pil_image)
                        
                        # 尝试加载中文字体
                        try:
                            font = ImageFont.truetype("simhei.ttf", 30)
                        except:
                            try:
                                font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 30)
                            except:
                                font = ImageFont.load_default()
                        
                        # 绘制多重行为文本
                        if behavior_classes:
                            behavior_text = "行为: " + " + ".join(behavior_classes)
                        else:
                            behavior_text = "行为: 正常驾驶"
                            
                        draw.text((10, 10), behavior_text, font=font, fill=(0, 255, 0))
                        
                        # 转换回OpenCV格式
                        annotated_frame = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
                        
                    except Exception as e:
                        print(f"使用PIL绘制中文失败: {e}")
                        # 回退到英文显示
                        if behavior_classes:
                            behavior_text = "Behavior: " + " + ".join(behavior_classes)
                        else:
                            behavior_text = "Behavior: Normal Driving"
                        cv2.putText(annotated_frame, behavior_text, 
                                   (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
                    
                    self.result_ready.emit(annotated_frame, detected_objects, " + ".join(behavior_classes) if behavior_classes else "正常驾驶")
                    
                cap.release()
            else:
                # 图像文件或摄像头
                if isinstance(self.source, str):
                    frame = cv2.imread(self.source)
                    if frame is not None:
                        results = self.model(frame, conf=self.conf_threshold, verbose=False)
                        detected_objects = []
                        for result in results:
                            boxes = result.boxes
                            if boxes is not None:
                                for box in boxes:
                                    cls = int(box.cls.item())
                                    conf = box.conf.item()
                                    detected_objects.append((cls, conf))
                        
                        behavior_classes = self.classify_behavior(detected_objects)
                        annotated_frame = results[0].plot()
                        
                        # 同样修复图像文件的中文显示
                        try:
                            from PIL import Image, ImageDraw, ImageFont
                            import numpy as np
                            
                            pil_image = Image.fromarray(cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB))
                            draw = ImageDraw.Draw(pil_image)
                            
                            try:
                                font = ImageFont.truetype("simhei.ttf", 30)
                            except:
                                try:
                                    font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 30)
                                except:
                                    font = ImageFont.load_default()
                            
                            if behavior_classes:
                                behavior_text = "行为: " + " + ".join(behavior_classes)
                            else:
                                behavior_text = "行为: 正常驾驶"
                            draw.text((10, 10), behavior_text, font=font, fill=(0, 255, 0))
                            annotated_frame = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
                            
                        except Exception as e:
                            print(f"使用PIL绘制中文失败: {e}")
                            if behavior_classes:
                                behavior_text = "Behavior: " + " + ".join(behavior_classes)
                            else:
                                behavior_text = "Behavior: Normal Driving"
                            cv2.putText(annotated_frame, behavior_text, 
                                       (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
                        
                        self.result_ready.emit(annotated_frame, detected_objects, " + ".join(behavior_classes) if behavior_classes else "正常驾驶")
        
        except Exception as e:
            self.error_occurred.emit(str(e))
        finally:
            self.finished.emit()
    
    def stop(self):
        self.running = False
    
    def classify_behavior(self, detected_objects):
        """根据检测到的对象分类驾驶行为 - 支持多重行为"""
        classes = [obj[0] for obj in detected_objects]
        confidences = [obj[1] for obj in detected_objects]
        
        # 定义类别映射 - 根据您的模型实际输出调整
        class_names = {
            0: 'smoking',
            1: 'phone', 
            2: 'closed_eye',
            3: 'open_eye',
            4: 'seatbelt',
            5: 'no_seatbelt'
        }
        
        detected_names = [class_names.get(cls, 'unknown') for cls in classes]
        print(f"检测到的对象: {detected_names}, 置信度: {confidences}")  # 调试信息
        
        # 行为分类逻辑 - 支持多重行为
        behavior_classes = []
        
        # 检查各种行为条件 - 使用较低的置信度阈值来确保检测到所有行为
        valid_objects = [(name, conf) for name, conf in zip(detected_names, confidences) if conf >= 0.2]
        valid_names = [name for name, conf in valid_objects]
        
        print(f"有效检测对象: {valid_names}")  # 调试信息
        
        # 疲劳驾驶检测
        if 'closed_eye' in valid_names:
            behavior_classes.append("疲劳驾驶")
            print("检测到疲劳驾驶条件: closed_eye")  # 调试信息
        
        # 分心驾驶检测
        if 'smoking' in valid_names:
            behavior_classes.append("分心驾驶")
            print("检测到分心驾驶条件: smoking")  # 调试信息
        
        if 'phone' in valid_names:
            behavior_classes.append("分心驾驶")
            print("检测到分心驾驶条件: phone")  # 调试信息
        
        # 危险驾驶检测
        if 'no_seatbelt' in valid_names:
            behavior_classes.append("危险驾驶")
            print("检测到危险驾驶条件: no_seatbelt")  # 调试信息
        
        # 去重
        behavior_classes = list(set(behavior_classes))
        
        # 如果检测到睁眼和系安全带，且没有其他危险行为，则认为是正常驾驶
        if ('open_eye' in valid_names and 'seatbelt' in valid_names and 
            not behavior_classes):
            behavior_classes.append("正常驾驶")
        
        # 如果没有检测到任何有效行为，但有检测结果
        if not behavior_classes and valid_objects:
            # 检查是否有睁眼但没系安全带等情况
            if 'open_eye' in valid_names and 'seatbelt' not in valid_names:
                behavior_classes.append("危险驾驶")
            else:
                behavior_classes.append("未知状态")
        
        print(f"最终行为分类: {behavior_classes}")  # 调试信息
        return behavior_classes

class FatigueDetectionApp(QMainWindow):
    def __init__(self):
        super().__init__()
        self.model = None
        self.detection_thread = None
        self.init_ui()
        
    def init_ui(self):
        self.setWindowTitle("驾驶行为检测系统 - 疲劳/分心/危险驾驶检测")
        self.setGeometry(100, 100, 1200, 800)
        
        # 设置中文字体
        self.set_chinese_font()
        
        # 中央部件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        
        # 主布局
        main_layout = QHBoxLayout()
        central_widget.setLayout(main_layout)
        
        # 左侧控制面板
        control_panel = self.create_control_panel()
        main_layout.addWidget(control_panel, 1)
        
        # 右侧显示区域
        display_panel = self.create_display_panel()
        main_layout.addWidget(display_panel, 2)
        
    def set_chinese_font(self):
        """设置中文字体"""
        font = QFont()
        font.setFamily("Microsoft YaHei")  # 微软雅黑
        font.setPointSize(9)
        QApplication.setFont(font)
        
    def create_control_panel(self):
        panel = QGroupBox("控制面板")
        layout = QVBoxLayout()
        
        # 模型选择
        model_group = QGroupBox("模型设置")
        model_layout = QVBoxLayout()
        
        self.model_path_edit = QTextEdit()
        self.model_path_edit.setMaximumHeight(60)
        self.model_path_edit.setPlaceholderText("模型文件路径...")
        model_layout.addWidget(QLabel("模型路径:"))
        model_layout.addWidget(self.model_path_edit)
        
        btn_layout = QHBoxLayout()
        self.load_model_btn = QPushButton("加载模型")
        self.load_model_btn.clicked.connect(self.load_model)
        btn_layout.addWidget(self.load_model_btn)
        
        self.browse_model_btn = QPushButton("浏览")
        self.browse_model_btn.clicked.connect(self.browse_model)
        btn_layout.addWidget(self.browse_model_btn)
        
        model_layout.addLayout(btn_layout)
        model_group.setLayout(model_layout)
        layout.addWidget(model_group)
        
        # 检测设置
        detect_group = QGroupBox("检测设置")
        detect_layout = QVBoxLayout()
        
        # 置信度阈值
        conf_layout = QHBoxLayout()
        conf_layout.addWidget(QLabel("置信度阈值:"))
        self.conf_spinbox = QDoubleSpinBox()
        self.conf_spinbox.setRange(0.1, 1.0)
        self.conf_spinbox.setValue(0.3)  # 降低默认阈值
        self.conf_spinbox.setSingleStep(0.05)
        conf_layout.addWidget(self.conf_spinbox)
        detect_layout.addLayout(conf_layout)
        
        # 输入源选择
        source_layout = QHBoxLayout()
        source_layout.addWidget(QLabel("输入源:"))
        self.source_combo = QComboBox()
        self.source_combo.addItems(["摄像头", "图像文件", "视频文件"])
        source_layout.addWidget(self.source_combo)
        detect_layout.addLayout(source_layout)
        
        detect_group.setLayout(detect_layout)
        layout.addWidget(detect_group)
        
        # 操作按钮
        button_group = QGroupBox("操作")
        button_layout = QVBoxLayout()
        
        self.start_btn = QPushButton("开始检测")
        self.start_btn.clicked.connect(self.start_detection)
        self.start_btn.setEnabled(False)
        button_layout.addWidget(self.start_btn)
        
        self.stop_btn = QPushButton("停止检测")
        self.stop_btn.clicked.connect(self.stop_detection)
        self.stop_btn.setEnabled(False)
        button_layout.addWidget(self.stop_btn)
        
        self.snapshot_btn = QPushButton("截图保存")
        self.snapshot_btn.clicked.connect(self.take_snapshot)
        self.snapshot_btn.setEnabled(False)
        button_layout.addWidget(self.snapshot_btn)
        
        button_group.setLayout(button_layout)
        layout.addWidget(button_group)
        
        # 检测结果
        result_group = QGroupBox("检测结果")
        result_layout = QVBoxLayout()
        
        self.behavior_label = QLabel("行为分类: 未检测")
        self.behavior_label.setStyleSheet("font-size: 14px; font-weight: bold; color: blue;")
        result_layout.addWidget(self.behavior_label)
        
        self.detection_text = QTextEdit()
        self.detection_text.setMaximumHeight(150)
        self.detection_text.setPlaceholderText("检测结果将显示在这里...")
        result_layout.addWidget(self.detection_text)
        
        result_group.setLayout(result_layout)
        layout.addWidget(result_group)
        
        # 类别说明
        info_group = QGroupBox("类别说明")
        info_layout = QVBoxLayout()
        
        info_text = """
        检测类别:
        • smoking - 吸烟
        • phone - 使用手机  
        • closed_eye - 闭眼
        • open_eye - 睁眼
        • seatbelt - 系安全带
        • no_seatbelt - 未系安全带
        
        行为分类:
        • 疲劳驾驶: 检测到闭眼
        • 分心驾驶: 检测到吸烟或使用手机
        • 危险驾驶: 检测到未系安全带
        • 正常驾驶: 检测到睁眼且系安全带
        • 多重行为: 可同时检测多种行为
        """
        info_label = QLabel(info_text)
        info_label.setStyleSheet("font-size: 12px;")
        info_layout.addWidget(info_label)
        
        info_group.setLayout(info_layout)
        layout.addWidget(info_group)
        
        layout.addStretch()
        panel.setLayout(layout)
        return panel
    
    def create_display_panel(self):
        panel = QGroupBox("视频显示")
        layout = QVBoxLayout()
        
        # 视频显示标签
        self.video_label = QLabel()
        self.video_label.setAlignment(Qt.AlignCenter)
        self.video_label.setMinimumSize(640, 480)
        self.video_label.setStyleSheet("border: 2px solid gray; background-color: black;")
        self.video_label.setText("视频显示区域")
        layout.addWidget(self.video_label)
        
        # 状态栏
        status_layout = QHBoxLayout()
        status_layout.addWidget(QLabel("状态:"))
        self.status_label = QLabel("就绪")
        self.status_label.setStyleSheet("color: green;")
        status_layout.addWidget(self.status_label)
        status_layout.addStretch()
        
        self.fps_label = QLabel("FPS: --")
        status_layout.addWidget(self.fps_label)
        
        layout.addLayout(status_layout)
        panel.setLayout(layout)
        return panel
    
    def load_model(self):
        model_path = self.model_path_edit.toPlainText().strip()
        if not model_path:
            QMessageBox.warning(self, "警告", "请输入模型文件路径")
            return
        
        if not os.path.exists(model_path):
            QMessageBox.critical(self, "错误", f"模型文件不存在: {model_path}")
            return
        
        try:
            self.model = YOLO(model_path)
            self.start_btn.setEnabled(True)
            self.status_label.setText("模型加载成功")
            self.status_label.setStyleSheet("color: green;")
            QMessageBox.information(self, "成功", "模型加载成功！")
        except Exception as e:
            QMessageBox.critical(self, "错误", f"模型加载失败: {str(e)}")
    
    def browse_model(self):
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择模型文件", "", "Model Files (*.pt *.pth)")
        if file_path:
            self.model_path_edit.setText(file_path)
    
    def start_detection(self):
        if self.model is None:
            QMessageBox.warning(self, "警告", "请先加载模型")
            return
        
        source_type = self.source_combo.currentText()
        source = 0  # 默认摄像头
        
        if source_type == "图像文件":
            file_path, _ = QFileDialog.getOpenFileName(
                self, "选择图像文件", "", "Image Files (*.jpg *.jpeg *.png *.bmp)")
            if not file_path:
                return
            source = file_path
        elif source_type == "视频文件":
            file_path, _ = QFileDialog.getOpenFileName(
                self, "选择视频文件", "", "Video Files (*.mp4 *.avi *.mov)")
            if not file_path:
                return
            source = file_path
        
        # 创建检测线程
        self.detection_thread = DetectionThread(
            self.model_path_edit.toPlainText().strip(),
            source,
            self.conf_spinbox.value()
        )
        
        self.detection_thread.result_ready.connect(self.update_detection_result)
        self.detection_thread.error_occurred.connect(self.handle_detection_error)
        self.detection_thread.finished.connect(self.detection_finished)
        
        self.start_btn.setEnabled(False)
        self.stop_btn.setEnabled(True)
        self.snapshot_btn.setEnabled(True)
        self.status_label.setText("检测中...")
        self.status_label.setStyleSheet("color: orange;")
        
        self.detection_thread.start()
    
    def stop_detection(self):
        if self.detection_thread and self.detection_thread.isRunning():
            self.detection_thread.stop()
            self.detection_thread.wait()
    
    def detection_finished(self):
        self.start_btn.setEnabled(True)
        self.stop_btn.setEnabled(False)
        self.snapshot_btn.setEnabled(False)
        self.status_label.setText("检测已停止")
        self.status_label.setStyleSheet("color: red;")
    
    def update_detection_result(self, frame, objects, behavior_class):
        # 转换图像格式用于显示
        rgb_image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        h, w, ch = rgb_image.shape
        bytes_per_line = ch * w
        qt_image = QImage(rgb_image.data, w, h, bytes_per_line, QImage.Format_RGB888)
        
        # 缩放图像以适应显示区域
        scaled_pixmap = QPixmap.fromImage(qt_image).scaled(
            self.video_label.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation)
        self.video_label.setPixmap(scaled_pixmap)
        
        # 更新行为分类 - 确保中文正确显示
        self.behavior_label.setText(f"行为分类: {behavior_class}")
        
        # 更新检测结果文本
        class_names = ['smoking', 'phone', 'closed_eye', 'open_eye', 'seatbelt', 'no_seatbelt']
        result_text = f"检测到 {len(objects)} 个对象:\n"
        for i, (cls, conf) in enumerate(objects):
            class_name = class_names[cls] if cls < len(class_names) else f'class_{cls}'
            result_text += f"{i+1}. {class_name}: {conf:.3f}\n"
        
        self.detection_text.setText(result_text)
        
        # 根据行为分类设置颜色 - 多重行为使用特殊颜色
        if " + " in behavior_class:
            color = "red"  # 多重行为用红色突出显示
        else:
            color_map = {
                "疲劳驾驶": "orange",
                "分心驾驶": "purple", 
                "危险驾驶": "red",
                "正常驾驶": "green",
                "未知状态": "gray",
                "未检测": "blue"
            }
            color = color_map.get(behavior_class, "black")
        
        self.behavior_label.setStyleSheet(f"font-size: 14px; font-weight: bold; color: {color};")
    
    def handle_detection_error(self, error_msg):
        QMessageBox.critical(self, "检测错误", error_msg)
        self.detection_finished()
    
    def take_snapshot(self):
        if hasattr(self, 'video_label') and self.video_label.pixmap():
            file_path, _ = QFileDialog.getSaveFileName(
                self, "保存截图", "", "PNG Files (*.png);;JPEG Files (*.jpg)")
            if file_path:
                self.video_label.pixmap().save(file_path)
                QMessageBox.information(self, "成功", f"截图已保存到: {file_path}")

def main():
    app = QApplication(sys.argv)
    
    # 设置应用程序样式
    app.setStyle('Fusion')
    
    window = FatigueDetectionApp()
    window.show()
    
    sys.exit(app.exec_())

if __name__ == '__main__':
    main()